import tempfile
import zipfile
import os
import magic

def spool_file_tomiddle(file):
    """
    将 SpooledTemporaryFile 的内容写入临时文件，并返回临时文件路径。
    :param file: SpooledTemporaryFile 对象
    :return: 临时文件路径
    """
    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
        file.seek(0)
        temp_file.write(file.read())
        temp_file_path = temp_file.name
    return temp_file_path

def get_file_extension(file_content):
    """
    根据文件内容检测文件类型并返回对应的后缀名。
    :param file_content: 文件内容（字节）
    :return: 文件后缀名（如 ".txt", ".png", ".mp4"）
    """
    mime = magic.Magic(mime=True)
    mime_type = mime.from_buffer(file_content)

    # 根据 MIME 类型返回对应的后缀名
    if mime_type == "text/plain":
        return ".txt"
    elif mime_type == "image/png":
        return ".png"
    elif mime_type == "image/jpeg":
        return ".jpg"
    elif mime_type == "application/pdf":
        return ".pdf"
    elif mime_type == "application/zip":
        return ".zip"
    elif mime_type == "video/mp4":
        return ".mp4"
    elif mime_type == "video/x-msvideo":
        return ".avi"
    elif mime_type == "video/x-matroska":
        return ".mkv"
    elif mime_type == "video/quicktime":
        return ".mov"
    elif mime_type == "application/octet-stream":
        # 检查文件头以识别视频格式
        if file_content.startswith(b'\x00\x00\x00\x18ftypmp42'):
            return ".mp4"
        elif file_content.startswith(b'\x52\x49\x46\x46'):  # RIFF header for AVI
            return ".avi"
        elif file_content.startswith(b'\x1A\x45\xDF\xA3'):  # Matroska header for MKV
            return ".mkv"
        elif file_content.startswith(b'\x00\x00\x00\x14ftypqt'):  # QuickTime header for MOV
            return ".mov"
    # 默认后缀名
    return ".bin"

def compress_files_to_zip(spooled_file: list, local_file_path, zip_file_path='./output/output.zip'):
    """
    将 SpooledTemporaryFile 对象和本地文件压缩到一个 ZIP 文件中。
    :param spooled_file: SpooledTemporaryFile 对象列表
    :param local_file_path: 本地文件的路径
    :param zip_file_path: 输出的 ZIP 文件路径，默认为 './output/output.zip'
    :return: None
    """
    temp_paths = []
    try:
        # 处理 SpooledTemporaryFile 对象
        for file in spooled_file:
            # 将 SpooledTemporaryFile 的内容写入临时文件
            temp_file_path = spool_file_tomiddle(file)
            # 读取文件内容以检测文件类型
            with open(temp_file_path, "rb") as f:
                file_content = f.read()
            # 获取文件后缀名
            extension = get_file_extension(file_content)
            # 重命名临时文件，添加后缀名
            new_temp_file_path = temp_file_path + extension
            os.rename(temp_file_path, new_temp_file_path)
            temp_paths.append(new_temp_file_path)

        # 创建 ZIP 文件
        with zipfile.ZipFile(zip_file_path, 'w') as zipf:
            # 将 SpooledTemporaryFile 的内容添加到 ZIP 文件中
            for temp_path in temp_paths:
                zipf.write(temp_path, os.path.basename(temp_path))
            # 将本地文件添加到 ZIP 文件中
            zipf.write(local_file_path, os.path.basename(local_file_path))
    finally:
        # 清理临时文件
        for temp_path in temp_paths:
            if os.path.exists(temp_path):
                os.remove(temp_path)
        # 关闭 SpooledTemporaryFile
        for file in spooled_file:
            file.close()

# 示例用法
if __name__ == "__main__":
    # 创建 SpooledTemporaryFile 对象
    spooled_file = []
    spooled_file1 = tempfile.SpooledTemporaryFile(max_size=1024)
    spooled_file1.write(b'Test1')
    spooled_file1.seek(0)  # 将文件指针重置到开头
    spooled_file2 = tempfile.SpooledTemporaryFile(max_size=1024)
    spooled_file2.write(b'Test2')
    spooled_file2.seek(0)  # 将文件指针重置到开头
    spooled_file.extend([spooled_file1, spooled_file2])

    # 本地文件路径
    local_file_path = '../docxs/output/example.docx'

    # 调用方法压缩文件
    compress_files_to_zip(spooled_file, local_file_path)

    print(f"ZIP 文件已生成：{os.path.abspath('./output/output.zip')}")